import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import cv2
import json
import numpy as np
from calibrate import CalibrationParams

class StereoCamera(object):
    def __init__(self):
        self.MatL = np.array([[1.7896234023549898e+03, 0.0, 1.2411994901617113e+03],
                            [0.0,1.7905028020864954e+03, 1.0599008022756107e+03],
                            [0.0, 0.0, 1]])
        self.DcoeffL = np.array([-5.5938406497188280e-02, 4.8075752564729526e-02, -3.4710198068299282e-04, -7.7254841935033242e-04, 1.8305947674082113e-03])

        self.MatR = np.array([[1.7765666219838608e+03, 0.0, 1.2411768619982781e+03], 
                              [0.0, 1.7771178345500975e+03, 1.0204299003542411e+03],
                              [0.0, 0.0, 1.0 ]])
        self.DcoeffR = np.array([ -6.9568008283865709e-02, 9.1313367474781013e-02, -2.1330075267639954e-04, 3.2860128024544971e-04, -4.8183486531339879e-02])

        self.R = np.array([[4.2913479156256279e-01, 6.7564041871490066e-02,-9.0070995937457210e-01],
                           [-7.4378860697889987e-02, 9.9645501192401498e-01, 3.9308959446862655e-02],
                           [9.0017282549066080e-01, 5.0124938478752816e-02, 4.3263885030202914e-01]])
        
        self.T = np.array([2.8221894388410453e+02, -1.1909325935800588e+01,1.5299481233902284e+02])

class Stereo(object):
    def __init__(self, calibration_param_path):
        with open(calibration_param_path, 'r') as f:
            json_dict = json.load(f)

        self.CameraMatL = np.array(json_dict['CameraMatL']) # 3*3
        self.DistCoeffL = np.array(json_dict['DistCoeffL']) # 1*5
        self.CameraMatR = np.array(json_dict['CameraMatR']) # 3*3
        self.DistCoeffR = np.array(json_dict['DistCoeffR']) # 1*5
        self.rvecs = np.array(json_dict['rvec']) # 38*3*1
        self.tvecs = np.array(json_dict['tvec']) # 38*3*1
        # self.WorldPoints = np.array(json_dict['objPoints']) # 38*88*3
        # self.imagePoints = np.array(json_dict['imagePoints']) # 38*88*1*2
        # self.imagePaths = json_dict['imagePaths']

class MonocularCamera(object):
    def __init__(self, calibration_param_path):
        with open(calibration_param_path, 'r') as f:
            json_dict = json.load(f)

        self.CameraMat = np.array(json_dict['CameraMat']) # 3*3
        self.DistCoeff = np.array(json_dict['DistCoeff']) # 1*5
        self.rvecs = np.array(json_dict['rvec']) # 38*3*1
        self.tvecs = np.array(json_dict['tvec']) # 38*3*1
        self.WorldPoints = np.array(json_dict['objPoints']) # 38*88*3
        self.imagePoints = np.array(json_dict['imagePoints']) # 38*88*1*2
        self.imagePaths = json_dict['imagePaths']

def revise_img(img_root, mtx, dist):
    mtx = np.array(mtx)
    dist = np.array(dist).reshape(1, 5)
    img = cv2.imread(img_root)
    h, w = img.shape[:2]
    newcameramtx, roi = cv2.getOptimalNewCameraMatrix(mtx, dist, (w, h), 1, (w, h))
    # undistort
    dst = cv2.undistort(img, mtx, dist, None, newcameramtx)
    # crop the image
    x, y, w, h = roi
    # dst = dst[y:y + h, x:x + w]
    cv2.imwrite('./revise_img_zhen.jpg', dst)
    return dst

def convert_world2camera(world_points, rvec, tvec):
    """
        世界坐标系 -> 相机坐标系
        camera_points = R @ world_points.T + tvec
        :return:
    """
    R, _ = cv2.Rodrigues(rvec)
    camera_points = R @ world_points.T + tvec

    # print(camera_point.shape)

    # fig = plt.figure()
    # ax = fig.add_subplot(111, projection='3d')

    # ax.scatter(camera_points[0,...], camera_points[1,...], camera_points[2,...], c='r', marker='o')

    # ax.set_xlabel('X')
    # ax.set_ylabel('Y')
    # ax.set_zlabel('Z')

    # plt.show()
    return camera_points

def convert_camera2world(camera_points, rvec, tvec):
        """
        相机坐标系 -> 世界坐标系
        world_points = R.T @ camera_points - R.T @ tvec
        :return:
        """
        R, _ = cv2.Rodrigues(rvec)
        world_points = R.T @ camera_points - R.T @ tvec

        # print(camera_point.shape)

        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')

        ax.scatter(world_points[0,...], world_points[1,...], world_points[2,...], c='r', marker='o')

        ax.set_xlabel('X')
        ax.set_ylabel('Y')
        ax.set_zlabel('Z')

        plt.show()

        return world_points

def convert_camera2pixel(camera_points, mat):
    """
    相机坐标系 -> 像素坐标系
    """
    # _, num = camera_points.shape
    # temp = np.ones(shape=(1, num))

    # new_camera_points = np.vstack((camera_points, temp))

    pixel_points = mat @ camera_points

    new_pixel_points = pixel_points[:2, :] / pixel_points[2, :]
 
    return new_pixel_points

# 消除畸变
def revise_img(image_path, revise_image_path, camera_mat, dist_coeff):
    camera_mat = np.array(camera_mat)
    dist_coeff = np.array(dist_coeff).reshape(1, 5)
    image_bgr = cv2.imread(image_path)
    h, w = image_bgr.shape[:2]
    newcameramtx, roi = cv2.getOptimalNewCameraMatrix(camera_mat, dist_coeff, (w, h), 1, (w, h))
    # undistort
    dst = cv2.undistort(image_bgr, camera_mat, dist_coeff, None, newcameramtx)
    # crop the image
    # x, y, w, h = roi
    # dst = dst[y:y + h, x:x + w]
    # cv2.imwrite(revise_image_path, dst)
    return dst

if __name__ == "__main__":

    cam_param = CalibrationParams()
    cam_param.load_monocular_param(r"D:\VSCodeProjects\Calibration\monocular_calibiation_L1210.json")
    
    camera_points = convert_world2camera(cam_param.worldPoints[0], cam_param.rvecs[0], cam_param.tvecs[0])
    
    pixel_points = convert_camera2pixel(camera_points, cam_param.CameraMat)
    print(pixel_points)

    # image = cv2.imread(cam_param.imagePaths[0])
    image = revise_img(cam_param.imagePaths[0], )
    for pt in pixel_points.T:
        cv2.circle(image, (int(pt[0]),int(pt[1])), 1, (0,0,255), 4)
        # img[pt] = 255
    cv2.imwrite('pixel.bmp', image)
    
    # cam = MonocularCamera(r"D:\VSCodeProjects\Calibration\monocular_calibiation1208R.json")
    # # stereo = StereoCamera()
    # stereo = Stereo(r"D:\VSCodeProjects\Calibration\stereo_calibiate1208.json")

    # revise_img(cam.imagePaths[0], cam.CameraMat, cam.DistCoeff)
    # revise_img(cam.imagePaths[0], stereo.MatR, stereo.DcoeffR)
    # revise_img(cam.imagePaths[0], stereo.CameraMatR, stereo.DistCoeffR)

    # world_points = np.zeros((88, 3), np.float32)
    # print(np.mgrid[0:11, 0:8])
    # world_points[:, :2] = np.mgrid[0:11, 0:8].T.reshape(-1, 2)
    # objp = world_points * 10
    # print(objp)
    # mesh_grid = np.mgrid[0:8:1, 0:8:1, 0:8:1]
    # # print(mesh_grid)

    # print(mesh_grid.shape)
    # print(mesh_grid[:, 0, 1, 0])
    
    # 
    # imgpts, _ = cv2.projectPoints(cam.WorldPoints[0], cam.rvecs[0], cam.tvecs[0], cam.CameraMat, cam.DistCoeff)
    # imgpts = np.int32(imgpts.reshape(-1, 2))

    # img = cv2.imread(r"D:\VSCodeProjects\Calibration\revise_img_zhen.jpg")
    

    
    

    # print(imgpts)
    
    # R 是从世界坐标系到相机坐标系的旋转矩阵
    # R.T (R的转置) 是从相机坐标系到世界坐标系的旋转矩阵
    # -R.T @ tvecs[0] 是从相机坐标系的原点（0,0,0）到世界坐标系的平移向量
    
    
